#!/usr/bin/env python3
"""
SSH Connection Manager for AI-Powered Terminal
Handles SSH connections using paramiko
"""

import asyncio
import logging
import socket
import threading
import time
from typing import Dict, List, Optional, Callable
from datetime import datetime

import paramiko
from fastapi import WebSocket

from models import ConnectionInfo
from database import DatabaseManager

# Configure more detailed logging for SSH operations
logging.basicConfig(level=logging.DEBUG)
logger = logging.getLogger(__name__)

class SSHSession:
    """Individual SSH session handler"""
    
    def __init__(self, session_id: str, connection_info: ConnectionInfo, websocket: WebSocket):
        self.session_id = session_id
        self.connection_info = connection_info
        self.websocket = websocket
        self.ssh_client: Optional[paramiko.SSHClient] = None
        self.channel: Optional[paramiko.Channel] = None
        self.connected = False
        self.output_thread: Optional[threading.Thread] = None
        self._stop_output = False
        # Store the current event loop for thread-safe operations
        try:
            self.loop = asyncio.get_running_loop()
        except RuntimeError:
            self.loop = None
            logger.warning(f"No running event loop found for session {session_id}")
    
    async def connect(self) -> bool:
        """Establish SSH connection"""
        logger.info(f"Starting SSH connection to {self.connection_info.hostname}:{self.connection_info.port}")
        try:
            self.ssh_client = paramiko.SSHClient()
            self.ssh_client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
            logger.debug(f"SSH client created for {self.session_id}")
            
            # Enable detailed SSH debugging
            paramiko.util.log_to_file('paramiko.log')
            
            # Connection parameters
            connect_params = {
                'hostname': self.connection_info.hostname,
                'port': self.connection_info.port,
                'username': self.connection_info.username,
                'timeout': 30,  # Increased timeout
                'auth_timeout': 30,
                'banner_timeout': 30
            }
            
            # Add authentication
            if self.connection_info.password:
                connect_params['password'] = self.connection_info.password
                logger.debug(f"Using password authentication for {self.session_id}")
            elif self.connection_info.key_file:
                connect_params['key_filename'] = self.connection_info.key_file
                logger.debug(f"Using key file authentication for {self.session_id}")
            else:
                logger.warning(f"No authentication method specified for {self.session_id}")
            
            # Test connectivity first
            logger.info(f"Testing network connectivity to {self.connection_info.hostname}:{self.connection_info.port}")
            try:
                test_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
                test_socket.settimeout(10)
                result = test_socket.connect_ex((self.connection_info.hostname, self.connection_info.port))
                test_socket.close()
                if result != 0:
                    logger.error(f"Network connectivity test failed to {self.connection_info.hostname}:{self.connection_info.port} - error code {result}")
                    await self._send_error(f"Cannot reach {self.connection_info.hostname}:{self.connection_info.port} - network error")
                    return False
                else:
                    logger.info(f"Network connectivity test successful to {self.connection_info.hostname}:{self.connection_info.port}")
            except Exception as e:
                logger.error(f"Network connectivity test error: {e}")
                await self._send_error(f"Network connectivity test failed: {str(e)}")
                return False
            
            # Connect to SSH server
            logger.debug(f"Attempting SSH connection with params: {dict(connect_params, password='***' if 'password' in connect_params else connect_params)}")
            self.ssh_client.connect(**connect_params)
            logger.info(f"SSH connection established to {self.connection_info.hostname}")
            
            # Create interactive shell channel
            logger.debug(f"Creating interactive shell channel for {self.session_id}")
            self.channel = self.ssh_client.invoke_shell(
                term='xterm-256color',
                width=80,
                height=24
            )
            logger.debug(f"Interactive shell channel created for {self.session_id}")
            
            self.connected = True
            
            # Start output monitoring thread
            self._start_output_monitoring()
            
            logger.info(f"SSH connection fully established: {self.session_id}")
            return True
            
        except paramiko.AuthenticationException as e:
            logger.error(f"SSH authentication failed for {self.session_id}: {e}")
            await self._send_error(f"Authentication failed: {str(e)}. Please check username and password.")
            return False
        except paramiko.SSHException as e:
            logger.error(f"SSH protocol error for {self.session_id}: {e}")
            await self._send_error(f"SSH protocol error: {str(e)}. Server may not support SSH or have connection limits.")
            return False
        except socket.timeout as e:
            logger.error(f"SSH connection timeout for {self.session_id}: {e}")
            await self._send_error(f"Connection timeout: {str(e)}. Server may be slow or unreachable.")
            return False
        except socket.error as e:
            logger.error(f"Network error for {self.session_id}: {e}")
            await self._send_error(f"Network error: {str(e)}. Check server address and port.")
            return False
        except Exception as e:
            logger.error(f"Unexpected SSH connection error for {self.session_id}: {e}", exc_info=True)
            await self._send_error(f"Connection failed: {str(e)}")
            return False
    
    def _start_output_monitoring(self):
        """Start thread to monitor SSH output"""
        logger.debug(f"Starting output monitoring thread for {self.session_id}")
        self.output_thread = threading.Thread(
            target=self._output_monitor,
            daemon=True
        )
        self.output_thread.start()
    
    def _output_monitor(self):
        """Monitor SSH output and send to WebSocket"""
        import asyncio
        try:
            while self.connected and not self._stop_output and self.channel:
                if self.channel.recv_ready():
                    try:
                        data = self.channel.recv(1024)
                        if data:
                            # Send output to WebSocket using thread-safe method
                            decoded_data = data.decode('utf-8', errors='ignore')
                            logger.debug(f"SSH output received from {self.session_id}: {len(decoded_data)} bytes")
                            self._send_output_sync(decoded_data)
                        else:
                            logger.debug(f"No more data from {self.session_id}, breaking")
                            break
                    except Exception as e:
                        logger.error(f"Error reading SSH output from {self.session_id}: {e}")
                        break
                
                time.sleep(0.01)  # Small delay to prevent CPU spinning
                
        except Exception as e:
            logger.error(f"Output monitor error for {self.session_id}: {e}")
        finally:
            if self.connected:
                logger.info(f"Output monitor ending for {self.session_id}")
                self._handle_disconnection_sync()
    
    def _send_output_sync(self, data: str):
        """Send output data to WebSocket (thread-safe)"""
        try:
            if self.loop and not self.loop.is_closed():
                # Use the stored event loop to schedule the coroutine from the thread
                asyncio.run_coroutine_threadsafe(self._send_output(data), self.loop)
            else:
                logger.warning(f"Cannot send output from thread - no valid event loop available for {self.session_id}")
        except Exception as e:
            logger.error(f"Failed to send output (sync): {e}")
    
    async def _send_output(self, data: str):
        """Send output data to WebSocket"""
        try:
            await self.websocket.send_json({
                "type": "terminal_output",
                "session_id": self.session_id,
                "data": data
            })
        except Exception as e:
            logger.error(f"Failed to send output: {e}")
    
    async def _send_error(self, error: str):
        """Send error message to WebSocket"""
        try:
            await self.websocket.send_json({
                "type": "ssh_error",
                "session_id": self.session_id,
                "error": error
            })
        except Exception as e:
            logger.error(f"Failed to send error: {e}")
    
    def _handle_disconnection_sync(self):
        """Handle SSH disconnection (thread-safe)"""
        self.connected = False
        try:
            if self.loop and not self.loop.is_closed():
                # Use the stored event loop to schedule the coroutine from the thread
                asyncio.run_coroutine_threadsafe(self._handle_disconnection(), self.loop)
            else:
                logger.warning(f"Cannot send disconnection notice from thread - no valid event loop available for {self.session_id}")
        except Exception as e:
            logger.error(f"Failed to handle disconnection (sync): {e}")
    
    async def _handle_disconnection(self):
        """Handle SSH disconnection"""
        self.connected = False
        try:
            await self.websocket.send_json({
                "type": "ssh_disconnected",
                "session_id": self.session_id,
                "message": "SSH connection lost"
            })
        except Exception as e:
            logger.error(f"Failed to send disconnection notice: {e}")
    
    async def send_input(self, data: str):
        """Send input to SSH channel"""
        if self.connected and self.channel:
            try:
                logger.debug(f"Sending input to {self.session_id}: {repr(data)}")
                self.channel.send(data)
            except Exception as e:
                logger.error(f"Failed to send input to {self.session_id}: {e}")
                await self._handle_disconnection()
        else:
            logger.warning(f"Cannot send input to {self.session_id}: not connected or no channel")
    
    async def resize_terminal(self, cols: int, rows: int):
        """Resize terminal"""
        if self.connected and self.channel:
            try:
                self.channel.resize_pty(width=cols, height=rows)
            except Exception as e:
                logger.error(f"Failed to resize terminal: {e}")
    
    async def execute_command(self, command: str) -> tuple[str, str]:
        """Execute a single command and return output"""
        if not self.connected or not self.ssh_client:
            raise Exception("SSH not connected")
        
        try:
            stdin, stdout, stderr = self.ssh_client.exec_command(command, timeout=30)
            
            # Read output
            output = stdout.read().decode('utf-8')
            error = stderr.read().decode('utf-8')
            
            return output, error
            
        except Exception as e:
            logger.error(f"Command execution failed: {e}")
            raise
    
    def disconnect(self):
        """Disconnect SSH session"""
        self.connected = False
        self._stop_output = True
        
        if self.output_thread and self.output_thread.is_alive():
            self.output_thread.join(timeout=1)
        
        if self.channel:
            try:
                self.channel.close()
            except:
                pass
        
        if self.ssh_client:
            try:
                self.ssh_client.close()
            except:
                pass
        
        logger.info(f"SSH session disconnected: {self.session_id}")

class SSHConnectionManager:
    """Manages multiple SSH connections"""
    
    def __init__(self):
        self.sessions: Dict[str, SSHSession] = {}
        self.client_sessions: Dict[str, list] = {}  # client_id -> list of session_ids
        self.db_manager = DatabaseManager()
    
    async def create_connection(self, client_id: str, connection_info: ConnectionInfo, websocket: WebSocket) -> str:
        """Create new SSH connection"""
        session_id = f"ssh_{client_id}_{int(time.time())}"
        logger.info(f"Creating SSH connection for client {client_id} to {connection_info.hostname}:{connection_info.port}")
        
        # Create SSH session
        session = SSHSession(session_id, connection_info, websocket)
        
        # Attempt connection
        if await session.connect():
            self.sessions[session_id] = session
            logger.info(f"SSH session {session_id} added to active sessions. Total: {len(self.sessions)}")
            
            # Track client sessions
            if client_id not in self.client_sessions:
                self.client_sessions[client_id] = []
            self.client_sessions[client_id].append(session_id)
            logger.debug(f"Client {client_id} now has {len(self.client_sessions[client_id])} sessions")
            
            # Update database usage
            # Note: We'd need to get the connection ID from the database
            # For now, we'll log the usage
            logger.info(f"SSH connection created successfully: {session_id} for {connection_info.hostname}")
            
            return session_id
        else:
            logger.error(f"Failed to establish SSH connection to {connection_info.hostname}")
            raise Exception(f"Failed to connect to {connection_info.hostname}")
    
    async def send_input(self, session_id: str, data: str):
        """Send input to SSH session"""
        if session_id in self.sessions:
            await self.sessions[session_id].send_input(data)
        else:
            logger.warning(f"Session not found: {session_id}")
    
    async def resize_terminal(self, session_id: str, cols: int, rows: int):
        """Resize terminal for session"""
        if session_id in self.sessions:
            await self.sessions[session_id].resize_terminal(cols, rows)
    
    async def execute_command(self, session_id: str, command: str) -> tuple[str, str]:
        """Execute command in SSH session"""
        if session_id in self.sessions:
            return await self.sessions[session_id].execute_command(command)
        else:
            raise Exception(f"Session not found: {session_id}")
    
    def disconnect_session(self, session_id: str):
        """Disconnect specific SSH session"""
        if session_id in self.sessions:
            session = self.sessions[session_id]
            session.disconnect()
            del self.sessions[session_id]
            
            # Remove from client sessions
            for client_id, session_list in self.client_sessions.items():
                if session_id in session_list:
                    session_list.remove(session_id)
                    break
    
    def cleanup_client_connections(self, client_id: str):
        """Clean up all connections for a client"""
        if client_id in self.client_sessions:
            session_ids = self.client_sessions[client_id].copy()
            for session_id in session_ids:
                self.disconnect_session(session_id)
            del self.client_sessions[client_id]
            logger.info(f"Cleaned up connections for client: {client_id}")
    
    def close_all_connections(self):
        """Close all SSH connections"""
        session_ids = list(self.sessions.keys())
        for session_id in session_ids:
            self.disconnect_session(session_id)
        
        self.client_sessions.clear()
        logger.info("All SSH connections closed")
    
    def get_session_info(self, session_id: str) -> Optional[Dict]:
        """Get session information"""
        if session_id in self.sessions:
            session = self.sessions[session_id]
            return {
                "session_id": session_id,
                "hostname": session.connection_info.hostname,
                "username": session.connection_info.username,
                "connected": session.connected,
                "enterprise": session.connection_info.enterprise
            }
        return None
    
    def get_active_sessions(self) -> List[Dict]:
        """Get all active sessions"""
        return [
            self.get_session_info(session_id) 
            for session_id in self.sessions.keys()
        ]

class SSHCommandAnalyzer:
    """Analyzes SSH commands for intelligent suggestions"""
    
    def __init__(self, db_manager: DatabaseManager):
        self.db_manager = db_manager
    
    def analyze_command(self, command: str, connection_info: ConnectionInfo) -> Dict:
        """Analyze command and provide insights"""
        analysis = {
            "command": command,
            "type": self._classify_command(command),
            "risk_level": self._assess_risk(command),
            "suggestions": [],
            "related_commands": []
        }
        
        # Add suggestions based on command type
        if analysis["type"] == "file_operation":
            analysis["suggestions"].append("Consider using 'ls -la' to verify file permissions")
        elif analysis["type"] == "system_modification":
            analysis["suggestions"].append("Make sure to backup before system modifications")
        elif analysis["type"] == "network":
            analysis["suggestions"].append("Check firewall rules if connection issues occur")
        
        # Get related commands from history
        analysis["related_commands"] = self._get_related_commands(command)
        
        return analysis
    
    def _classify_command(self, command: str) -> str:
        """Classify command type"""
        cmd_lower = command.lower().strip()
        
        file_ops = ['ls', 'cp', 'mv', 'rm', 'mkdir', 'rmdir', 'chmod', 'chown']
        system_ops = ['systemctl', 'service', 'apt', 'yum', 'dnf', 'pip']
        network_ops = ['curl', 'wget', 'ping', 'telnet', 'ssh', 'scp']
        monitor_ops = ['top', 'htop', 'ps', 'df', 'free', 'iostat']
        
        first_word = cmd_lower.split()[0] if cmd_lower.split() else ""
        
        if first_word in file_ops:
            return "file_operation"
        elif first_word in system_ops:
            return "system_modification"
        elif first_word in network_ops:
            return "network"
        elif first_word in monitor_ops:
            return "monitoring"
        else:
            return "general"
    
    def _assess_risk(self, command: str) -> str:
        """Assess command risk level"""
        cmd_lower = command.lower()
        
        high_risk = ['rm -rf', 'mkfs', 'fdisk', 'dd if=', 'sudo su', 'chmod 777']
        medium_risk = ['rm ', 'mv ', 'systemctl stop', 'service stop', 'iptables']
        
        if any(risk in cmd_lower for risk in high_risk):
            return "high"
        elif any(risk in cmd_lower for risk in medium_risk):
            return "medium"
        else:
            return "low"
    
    def _get_related_commands(self, command: str) -> List[str]:
        """Get related commands from history"""
        # This would query the database for similar commands
        # For now, return some static related commands
        cmd_lower = command.lower()
        
        if 'ls' in cmd_lower:
            return ['ls -la', 'ls -lh', 'find . -name "*"', 'tree']
        elif 'ps' in cmd_lower:
            return ['ps aux', 'top', 'htop', 'systemctl status']
        elif 'docker' in cmd_lower:
            return ['docker ps', 'docker images', 'docker logs', 'docker exec -it']
        else:
            return []